# -*- coding:utf-8 -*-
"""
作者：ChenTao
日期：2023年03月12日

此文件包含主要内容：
    object_detect定义目标检测代码
    conveying_message_to_control取出队列数据，打包数据到主控板

注意：此多进程程序运行时，一个进程出问题会导致所有进程卡死。
    因为我把所有目标检测进程的数据汇总在一个串口进程，目标检测某一个进程卡死，会导致串口进程卡死，从而导致其他进程也卡死
    一般都是硬件连接会出问题，硬件没问题程序就可以跑起来
"""
import multiprocessing
import time

import yaml
import argparse
from util import Yolov5, Control_board, stop_message, GpioControl
from yolov5.detect import queen_rw
from camera.record_video import Record_Video

# ---------------------------------------- #
# 加载超config的参数
# ---------------------------------------- #
hym_arg = argparse.ArgumentParser()
hym_arg.add_argument('--yaml', type=str, default='config/pc.yaml', help=argparse.SUPPRESS)
hym_opt = hym_arg.parse_args()
with open(hym_opt.yaml, 'r', encoding='utf-8') as f:
    hyp = yaml.load(f.read(), Loader=yaml.SafeLoader)

# ---------------------------------------- #
# 一些全局变量
# ---------------------------------------- #
"""进程间通信数据"""
message = {'coordinate': None, 'angle': None}  # 用于放入队列的数据


def conveying_message_to_control(hyp, queen, queen_stop):
    """
    给主控板发消息
    Args:
        queen_stop:
        queen:
        hyp:

    Returns:

    """
    if 'raspberry' in hym_opt.yaml:
        gpiocontrol = GpioControl(18)
    control_board = {}
    is_stop = {}
    is_self_test = {}
    for key in hyp.keys():
        if hyp[key]['open']:
            control_board[key] = Control_board(port=hyp[key]['serial']['mpu_com_port'],
                                               baudrate=hyp[key]['serial']['baudrate'])  # 初始化与主控板的串口
    try:
        while True:
            for key in control_board.keys():
                if hyp[key]['open']:
                    is_stop[key], is_self_test[key] = stop_message(control_board[key], key)
                    # queen_stop[key].put(dict(is_stop=is_stop[key], is_self_test=is_self_test[key]))

                    # data = queen[key].get()
                    data = queen_rw(queen_r=queen[key], queen_w=queen_stop[key],
                                    w_data=dict(is_stop=is_stop[key], is_self_test=is_self_test[key]))
                    if not data:
                        print(f"Process: {'%s %s' % (key, multiprocessing.current_process().name):-<20}waiting......")
                        continue

                    try:
                        if data['coordinate'] is not None and hyp[key]['message']['coordinate']:  # 取出坐标数据
                            # 计算与相机中心点的坐标差
                            dcoordinate = control_board[key].coordinate_difference(data_tuple=data['coordinate'],
                                                                                   translation_vector=
                                                                                   hyp[key]['capture']['camera'][
                                                                                       'translation_vector'],
                                                                                   capture_size=
                                                                                   hyp[key]['capture']['camera'][
                                                                                       'size'])
                            # print('坐标差{}'.format(dcoordinate))
                            control_board[key].send_message(dcoordinate, key, pad_len=6)
                        if data['angle'] is not None and hyp[key]['message']['angle']:  # 取出角度数据
                            # print('角度{}'.format(data['angle']))
                            control_board[key].send_message(data['angle'], key, pad_len=6)  # 串口发送
                        if data['self_test'] is not None:
                            print(
                                f"Process: {'%s %s' % (key, multiprocessing.current_process().name):-<20}send message: {data['self_test']}")
                            control_board[key].mpu.send(send_data=data['self_test'])

                    # 串口断连，重新连接
                    except IOError:
                        if control_board[key]:
                            control_board[key].mpu.com_close()
                        control_board[key] = Control_board(port=hyp[key]['serial']['mpu_com_port'],
                                                           baudrate=hyp[key]['serial']['baudrate'])  # 初始化与主控板的串口

            if 'raspberry' in hym_opt.yaml:
                if is_stop['cirque']:
                    gpiocontrol.output_high()
                else:
                    gpiocontrol.output_low()

                try:
                    if is_stop['pvc'] and is_stop['cirque']:
                        gpiocontrol.flicker(0.01)
                    elif is_stop['pvc']:
                        gpiocontrol.output_high()
                    elif is_stop['cirque']:
                        gpiocontrol.output_low()
                except KeyError:
                    pass

    except KeyboardInterrupt:
        for key in control_board.keys():
            if control_board[key]:
                control_board[key].mpu.com_close()
        if 'raspberry' in hym_opt.yaml:
            gpiocontrol.gpio_close()
        print("程序终止")


def object_detect(hyp, queen, queen_stop):
    """目标检测函数"""

    if hyp['capture']['video'] == 'D435i' or hyp['capture']['video'] == 'd435i':
        """D435摄像头"""
        from yolov5.rstest import Run_D435i_With_YoloV5
        run_d435i_with_yolov5 = Run_D435i_With_YoloV5(weight=hyp['yolov5']['weight'],
                                                      data=hyp['yolov5']['data'],
                                                      img_size=hyp['yolov5']['img_size'],
                                                      view_img=hyp['yolov5']['view_img'],
                                                      conf_thres=hyp['yolov5']['threshold']['confidence'],
                                                      iou_thres=hyp['yolov5']['threshold']['iou'],
                                                      device=hyp['yolov5']['device'],
                                                      depth_size=hyp['capture']['d435i']['depth_size'],
                                                      color_size=hyp['capture']['d435i']['color_size'],
                                                      fps=hyp['capture']['d435i']['fps'])
        # 开启检测
        run_d435i_with_yolov5.run(queen, translation_vector=hyp['capture']['d435i']['translation_vector'])
    else:
        """普通单目摄像头"""
        cnn = Yolov5(weights_path=hyp['yolov5']['weight'],
                     source=hyp['capture']['video'],
                     data=hyp['yolov5']['data'],
                     img_size=hyp['yolov5']['img_size'],
                     cap_info=hyp['capture']['camera'],
                     view_img=hyp['yolov5']['view_img'],
                     conf_thres=hyp['yolov5']['threshold']['confidence'],
                     iou_thres=hyp['yolov5']['threshold']['iou'],
                     nosave=hyp['yolov5']['nosave'],
                     dnn=hyp['yolov5']['dnn'],
                     camera_matrix=hyp['capture']['camera']['camera_matrix'],
                     dist_coeffs=hyp['capture']['camera']['dist_coeffs'], )  # 初始化神经网络

        # 开启检测
        cnn.run_detect(queen, queen_stop)


def run__queen(funs):
    """开启多个进程"""

    from multiprocessing import Process, Queue, cpu_count

    # 队列实例化
    queen = {}
    queen_stop = {}
    for key in hyp.keys():
        if hyp[key]['open']:
            queen[key] = Queue(maxsize=8)
            queen_stop[key] = Queue(maxsize=8)

    # 定义多个进程，name：进程名，target：函数名，args：函数参数
    process = []
    for fun in funs:
        if 'Record_Video' in str(fun):
            for key in hyp.keys():
                process.append(Process(name=f'{key}-record_video', target=fun,
                                       args=(key, hyp[key]['capture']['video'], hyp[key]['capture']['camera']['size'],
                                             hyp[key]['capture']['camera']['fps'], 'camera/video')))

        else:
            if 'object_detect' in str(fun):
                for key in hyp.keys():
                    if hyp[key]['open']:
                        process.append(
                            Process(name=f"{key} detect", target=fun, args=(hyp[key], queen[key], queen_stop[key])))
            else:
                process.append(Process(name='serial', target=fun, args=(hyp, queen, queen_stop)))

    assert len(process) <= int(cpu_count()), '进程数多余cpu核数，请减少进程'

    # 开启多进程
    [p.start() for p in process]
    [p.join() for p in process]


if __name__ == "__main__":
    # 多进程，通过queen队列通信
    funs = [object_detect,
            conveying_message_to_control,
            # Record_Video
            ]

    print(str(funs))
    run__queen(funs)
